import math
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np


def focal_loss(input_values, gamma):
    """Computes the focal loss"""
    p = torch.exp(-input_values)
    loss = (1 - p) ** gamma * input_values
    return loss.mean()


class FocalLoss(nn.Module):
    def __init__(self, weight=None, gamma=0.):
        super(FocalLoss, self).__init__()
        assert gamma >= 0
        self.gamma = gamma
        self.weight = weight

    def forward(self, input, target):
        return focal_loss(F.cross_entropy(input, target, reduction='none', weight=self.weight), self.gamma)


class LDAMLoss(nn.Module):

    def __init__(self, cls_num_list, max_m=0.5, weight=None, s=30):
        super(LDAMLoss, self).__init__()
        m_list = 1.0 / np.sqrt(np.sqrt(cls_num_list))
        m_list = m_list * (max_m / np.max(m_list))
        m_list = torch.cuda.FloatTensor(m_list)
        self.m_list = m_list
        assert s > 0
        self.s = s
        self.weight = weight

    def forward(self, x, target):
        index = torch.zeros_like(x, dtype=torch.uint8)
        index.scatter_(1, target.data.view(-1, 1), 1)

        index_float = index.type(torch.cuda.FloatTensor)
        batch_m = torch.matmul(self.m_list[None, :], index_float.transpose(0, 1))
        batch_m = batch_m.view((-1, 1))
        x_m = x - batch_m

        output = torch.where(index, x_m, x)
        return F.cross_entropy(self.s * output, target, weight=self.weight)


class PaCoLoss(nn.Module):
    def __init__(self, alpha, beta=1.0, gamma=1.0, supt=1.0, temperature=1.0, base_temperature=None, K=128,
                 num_classes=1000):
        super(PaCoLoss, self).__init__()
        self.temperature = temperature
        self.base_temperature = temperature if base_temperature is None else base_temperature
        self.K = K
        self.alpha = alpha
        self.beta = beta
        self.gamma = gamma
        self.supt = supt
        self.num_classes = num_classes

    def cal_weight_for_classes(self, cls_num_list):
        cls_num_list = torch.Tensor(cls_num_list).view(1, self.num_classes)
        self.weight = cls_num_list / cls_num_list.sum()
        self.weight = self.weight.to(torch.device('cuda'))

    def forward(self, features, labels=None, sup_logits=None, mask=None, epoch=None):
        device = (torch.device('cuda')
                  if features.is_cuda
                  else torch.device('cpu'))

        ss = features.shape[0]
        batch_size = (features.shape[0] - self.K) // 2

        labels = labels.contiguous().view(-1, 1)
        mask = torch.eq(labels[:batch_size], labels.T).float().to(device)

        # compute logits
        anchor_dot_contrast = torch.div(
            torch.matmul(features[:batch_size], features.T),
            self.temperature)

        # add supervised logits
        anchor_dot_contrast = torch.cat(((sup_logits + torch.log(self.weight + 1e-9)) / self.supt, anchor_dot_contrast),
                                        dim=1)

        # for numerical stability
        logits_max, _ = torch.max(anchor_dot_contrast, dim=1, keepdim=True)
        logits = anchor_dot_contrast - logits_max.detach()

        # mask-out self-contrast cases
        logits_mask = torch.scatter(
            torch.ones_like(mask),
            1,
            torch.arange(batch_size).view(-1, 1).to(device),
            0
        )
        mask = mask * logits_mask

        # add ground truth
        one_hot_label = torch.nn.functional.one_hot(labels[:batch_size, ].view(-1, ), num_classes=self.num_classes).to(
            torch.float32)
        mask = torch.cat((one_hot_label * self.beta, mask * self.alpha), dim=1)

        # compute log_prob
        logits_mask = torch.cat((torch.ones(batch_size, self.num_classes).to(device), self.gamma * logits_mask), dim=1)
        exp_logits = torch.exp(logits) * logits_mask
        log_prob = logits - torch.log(exp_logits.sum(1, keepdim=True) + 1e-12)

        # compute mean of log-likelihood over positive
        mean_log_prob_pos = (mask * log_prob).sum(1) / mask.sum(1)

        # loss
        loss = - (self.temperature / self.base_temperature) * mean_log_prob_pos
        loss = loss.mean()

        return loss


class BSCE(nn.Module):
    def __init__(self, cls_num_list, device) -> None:
        super().__init__()
        self.device = device
        self.cls_num_list = torch.Tensor(cls_num_list).to(device)

    def forward(self, outputs, labels):
        spc = self.cls_num_list.type_as(outputs)
        spc = spc.unsqueeze(0).expand(outputs.shape[0], -1)
        outputs = outputs + spc.log()
        loss = F.cross_entropy(input=outputs, target=labels, reduction="mean").to(self.device)
        return loss


class GML(nn.Module):
    def __init__(self, device, num_class_list):
        super().__init__()
        self.p = 1
        self.device = device
        self.weight = torch.Tensor(num_class_list).to(device)

    def forward(self, output, target):
        assert len(target.size()) == 1  # Target should be of 1-Dim

        max_logit = torch.max(output, dim=1, keepdim=True)[0]  # of shape N x 1
        max_logit = max_logit.detach()
        logits = output - max_logit
        exp_logits = torch.exp(logits) * self.weight.view(-1, self.weight.shape[0])
        prob = torch.clamp(exp_logits / exp_logits.sum(1, keepdim=True), min=1e-5, max=1.)

        num_images, num_classes = prob.size()

        ground_class_prob = torch.gather(prob, dim=1, index=target.view(-1, 1))
        ground_class_prob = ground_class_prob.repeat((1, num_classes))

        mask = torch.zeros((num_images, num_classes), dtype=torch.int64, device=self.device)
        mask[range(num_images), target] = 1
        num_images_per_class = torch.sum(mask, dim=0)
        exist_class_mask = torch.zeros((num_classes,), dtype=torch.int64, device=self.device)
        exist_class_mask[num_images_per_class != 0] = 1

        num_images_per_class[num_images_per_class == 0] = 1  # avoid the dividing by zero exception

        mean_prob_classes = torch.sum(ground_class_prob * mask, dim=0) / num_images_per_class  # of shape (C,)
        mean_prob_classes[exist_class_mask == 1] = -torch.log(mean_prob_classes[exist_class_mask == 1])

        mean_prob_sum = torch.sum(torch.pow(mean_prob_classes[exist_class_mask == 1], self.p)) / torch.sum(
            exist_class_mask)

        loss = torch.pow(mean_prob_sum, 1.0 / self.p)

        return loss
